import psycopg2
from docx import Document
from docx.shared import Inches, RGBColor
import matplotlib.pyplot as plt
import networkx as nx

# 数据库配置
DB_CONFIG = {
    "host": "127.0.0.1",
    "dbname": "studydb",
    "user": "postgres",
    "password": "159230@Txl",
    "port": "5432"
}
TARGET_SCHEMA = "studys"

# 排除表的特征
EXCLUDE_PATTERNS = [
    lambda x: any(char.isdigit() for char in x),  # 表名中含有数字
    lambda x: x.endswith('_back') or x.endswith('_prod') or x.endswith('_copy'),  # 以 _back、_prod、_copy 结尾
    lambda x: x.startswith('temp_')  # 以 temp_ 开头
]


def should_exclude_table(table_name):
    """判断表是否需要排除"""
    return any(pattern(table_name) for pattern in EXCLUDE_PATTERNS)


def get_tables_and_columns(conn, schema):
    """获取指定模式下的所有表和字段信息，包括字段注释和表注释"""
    cursor = conn.cursor()

    # 获取表和字段信息
    cursor.execute(f"""
        SELECT c.table_name, c.column_name, c.data_type, c.is_nullable, c.column_default, c.ordinal_position, pgd.description
        FROM information_schema.columns c
        LEFT JOIN pg_catalog.pg_statio_all_tables st ON (c.table_schema = st.schemaname AND c.table_name = st.relname)
        LEFT JOIN pg_catalog.pg_description pgd ON (st.relid = pgd.objoid AND c.ordinal_position = pgd.objsubid)
        WHERE c.table_schema = '{schema}'
        ORDER BY c.table_name, c.ordinal_position;
    """)
    tables = {}
    for table_name, column_name, data_type, is_nullable, column_default, ordinal_position, description in cursor.fetchall():
        if table_name not in tables:
            tables[table_name] = {"comment": None, "columns": []}
        tables[table_name]["columns"].append({
            "column_name": column_name,
            "data_type": data_type,
            "is_nullable": is_nullable,
            "column_default": column_default,
            "ordinal_position": ordinal_position,
            "comment": description if description else "\\"  # 如果没有注释，填充为\
        })

    # 获取表的中文注释
    cursor.execute(f"""
        SELECT relname AS table_name, obj_description(oid) AS table_comment
        FROM pg_class
        WHERE relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = '{schema}')
        AND relkind = 'r';
    """)
    for table_name, table_comment in cursor.fetchall():
        if table_name in tables:
            tables[table_name]["comment"] = table_comment if table_comment else "无注释"

    cursor.close()
    return tables


def generate_er_diagram(tables, table_name, columns):
    """生成ER图并保存为图片"""
    G = nx.DiGraph()

    # 添加表名节点（长方形，白色填充，黑色边框）
    G.add_node(table_name, shape='box', color='black', style='filled', fillcolor='white')

    # 添加字段节点（椭圆，白色填充，黑色边框）
    for column in columns:
        G.add_node(column['column_name'], shape='ellipse', color='black', style='filled', fillcolor='white')
        G.add_edge(table_name, column['column_name'])

    # 绘制图形
    pos = nx.spring_layout(G)
    plt.figure(figsize=(10, 6))

    # 绘制表名节点（长方形）
    nx.draw_networkx_nodes(
        G, pos, nodelist=[table_name],
        node_shape='s',  # 长方形
        node_color='white',  # 填充颜色
        edgecolors='black',  # 边框颜色
        linewidths=1.5,  # 边框宽度
        node_size=5000  # 节点大小
    )

    # 绘制字段节点（椭圆）
    nx.draw_networkx_nodes(
        G, pos, nodelist=[col['column_name'] for col in columns],
        node_shape='o',  # 椭圆
        node_color='white',  # 填充颜色
        edgecolors='black',  # 边框颜色
        linewidths=1.5,  # 边框宽度
        node_size=3000  # 节点大小
    )

    # 绘制边
    nx.draw_networkx_edges(G, pos, arrowstyle='->', arrowsize=10, edge_color='black')

    # 绘制标签
    nx.draw_networkx_labels(G, pos, font_size=8, font_color='black')

    # 保存图片
    plt.axis('off')
    plt.savefig(f"{table_name}_er.png", format="PNG", bbox_inches='tight')
    plt.close()


def set_heading_style(heading, color):
    """设置标题样式（颜色和加粗）"""
    for run in heading.runs:
        run.font.color.rgb = color  # 设置颜色
        run.font.bold = False  # 不加粗


def create_word_document(tables):
    """生成Word文档"""
    doc = Document()
    doc.add_heading('XXXX系统数据库表设计文档', 0)

    # 设置主标题样式
    set_heading_style(doc.paragraphs[-1], RGBColor(0, 0, 0))

    # 1. 引言
    heading = doc.add_heading('1. 引言', level=1)
    set_heading_style(heading, RGBColor(0, 0, 0))
    doc.add_paragraph(
        '目的：本文档旨在详细描述XXXX系统的数据库表结构，包括每个数据表的名称、字段及字段类型、字段含义、相关枚举值含义以及数据库ER图。')
    doc.add_paragraph('范围：涵盖整个系统的数据库设计。')
    doc.add_paragraph('定义与术语：列出文档中使用的专业术语和缩写词。')

    # 2. 数据库概述
    heading = doc.add_heading('2. 数据库概述', level=1)
    set_heading_style(heading, RGBColor(0, 0, 0))
    doc.add_paragraph('数据库目标：简述数据库的设计目标和预期效果。')
    doc.add_paragraph('数据库架构：描述数据库的高层架构，包括主要的数据表和关系。')

    # 3. 数据表详细说明
    heading = doc.add_heading('3. 数据表详细说明', level=1)
    set_heading_style(heading, RGBColor(0, 0, 0))
    for i, (table_name, table_info) in enumerate(tables.items(), start=1):
        table_comment = table_info["comment"]  # 表的中文注释
        columns = table_info["columns"]

        # 添加标题：数据表中文注释（表名）
        heading = doc.add_heading(f'3.{i} {table_comment}（{table_name}）', level=2)
        set_heading_style(heading, RGBColor(0, 0, 0))
        doc.add_paragraph(f'表名称：{table_name}')
        table = doc.add_table(rows=1, cols=5)
        table.style = 'Table Grid'
        hdr_cells = table.rows[0].cells
        hdr_cells[0].text = '序号'
        hdr_cells[1].text = '字段'
        hdr_cells[2].text = '字段类型'
        hdr_cells[3].text = '字段含义'
        hdr_cells[4].text = '枚举值含义'
        for column in columns:
            row_cells = table.add_row().cells
            row_cells[0].text = str(column['ordinal_position'])
            row_cells[1].text = column['column_name']
            row_cells[2].text = column['data_type']
            row_cells[3].text = column['comment']  # 使用字段的中文注释
            row_cells[4].text = '\\'

    # 4. 数据库ER图
    heading = doc.add_heading('4. 数据库ER图', level=1)
    set_heading_style(heading, RGBColor(0, 0, 0))
    doc.add_paragraph('ER图说明：提供数据库的实体关系图，展示各个数据表之间的关系。')
    for table_name, table_info in tables.items():
        generate_er_diagram(tables, table_name, table_info["columns"])
        doc.add_picture(f"{table_name}_er.png", width=Inches(6))
        doc.add_paragraph()

    # 保存文档
    doc.save('XXXX系统数据库表设计文档.docx')


def main():
    # 连接数据库
    conn = psycopg2.connect(**DB_CONFIG)
    tables = get_tables_and_columns(conn, TARGET_SCHEMA)

    # 过滤掉不需要的表
    tables = {table_name: table_info for table_name, table_info in tables.items() if not should_exclude_table(table_name)}

    # 生成Word文档
    create_word_document(tables)

    # 关闭数据库连接
    conn.close()


if __name__ == "__main__":
    main()